#include<cmath>
#include<algorithm>
#include"system_m_hf.h"
#include"help.h"
using std::abs;
using std::min;
using std::max;

void System_TzM_HF::setupOrbitals()
{
  hfsolver.init();
  hfsolver.iter();
  hfsolver.calEnergy();
  HFenergy=hfsolver.Energy;
  //  Orbitals=hfsolver.HFOrbitals;
  //proton orbitals ahead, to use in shell model
  int totalOrbitals=hfsolver.HFOrbitals.size();
  Orbitals.resize(totalOrbitals);
  maxProtons=0;
  maxNeutrons=0;
  int j=0;
  for(int i=0;i<totalOrbitals;i++)
    {
      if(hfsolver.HFOrbitals[i].tz==-1)
	{
	  Orbitals[j]=hfsolver.HFOrbitals[i];
	  ++j;
	  ++maxProtons;
	}
    }
  for(int i=0;i<totalOrbitals;i++)
    {
      if(hfsolver.HFOrbitals[i].tz==1)
	{
	  Orbitals[j]=hfsolver.HFOrbitals[i];
	  ++j;
	  ++maxNeutrons;
	}
    }

  setupGroups();
  // //set fermi surface
  // FermiSurface=0;
  // for(int i=0;i<Orbitals.size();i++)
  //   {
  //     if(Orbitals[i].state) break;
  //     FermiSurface=i;
  //   }
}

/*
  dim of MatEle=4,the 4 double are:
      0  interaction
      1  (p1 \dot p2)/(m hbar omega) + m omega r1 \dot r2 / hbar
      2  (r1-r2)^2 m omega/(2 hbar)
      3  (p1 \dot p2)/(m hbar omega)
*/
const int System_TzM_HF::MatEledim=System_TzType::MatEledim;
void System_TzM_HF::setupTwoBodyMat()
{
  setupIndiceIn2BStates();
  int totalChannels=TwoBodyChannels.size();
  TwoBodyMat.resize(totalChannels);
  #pragma omp parallel for
  for(int channel=0;channel<totalChannels;channel++)
    {
      int num2BInChannel=TwoBodyChannels[channel].size();
      TwoBodyMat[channel].clear();
      TwoBodyMat[channel].resize( (num2BInChannel*(num2BInChannel+1))/2 );
      for(int bra=0;bra<num2BInChannel;bra++)
	for(int ket=bra;ket<num2BInChannel;ket++)
	  {
	    int ab=TwoBodyChannels[channel][bra];
	    int cd=TwoBodyChannels[channel][ket];
	    int a=TwoBodyStates[ab].a;
	    int b=TwoBodyStates[ab].b;
	    int c=TwoBodyStates[cd].a;
	    int d=TwoBodyStates[cd].b;
	    TwoBodyMatEle MatEle(MatEledim);
	    get2BmatOnNeed(a,b,c,d,MatEle);
	    set2Bmat(channel,bra,ket,MatEle);
	  }
    }
}

void System_TzM_HF::get2BmatOnNeed(int a,int b,int c,int d,TwoBodyMatEle & MatEle) const
{
  MatEle.resize(MatEledim);
  MatEle.setZero();

  if( (Orbitals[a].par+Orbitals[b].par+Orbitals[c].par+Orbitals[d].par)%2 ) return;
  if( (Orbitals[a].tz+Orbitals[b].tz)!=(Orbitals[c].tz+Orbitals[d].tz) ) return;
  if( (Orbitals[a].mm+Orbitals[b].mm)!=(Orbitals[c].mm+Orbitals[d].mm) ) return;  

  for(int ia=0;ia<Orbitals[a].Coeff.size();ia++)
    for(int ib=0;ib<Orbitals[b].Coeff.size();ib++)
      for(int ic=0;ic<Orbitals[c].Coeff.size();ic++)
	for(int id=0;id<Orbitals[d].Coeff.size();id++)
	  {
	    System_TzType::TwoBodyMatEle Temp;
	    int iia=pSystem_Tz->Groups[Orbitals[a].GroupIndex][ia];
	    int iib=pSystem_Tz->Groups[Orbitals[b].GroupIndex][ib];
	    int iic=pSystem_Tz->Groups[Orbitals[c].GroupIndex][ic];
	    int iid=pSystem_Tz->Groups[Orbitals[d].GroupIndex][id];
	    //pSystem_Tz->get2BmatAt(iia,iib,iic,iid,Temp);
	    pSystem_Tz->get2BmatOnNeed(iia,iib,iic,iid,Temp);
	    if(Temp.empty()) continue;
	    MatEle+=Temp*(Orbitals[a].Coeff[ia]*Orbitals[b].Coeff[ib]*Orbitals[c].Coeff[ic]*Orbitals[d].Coeff[id]);
	  }
}

System_TzM_HF::DataType System_TzM_HF::get1B(int bra,int ket) const
{
  DataType val=0;
  if(Orbitals[bra].isInAGroupWith(Orbitals[ket]) )
    {
      // for(int ibra=0;ibra<Orbitals[bra].Coeff.size();ibra++)
      // 	for(int iket=0;iket<Orbitals[ket].Coeff.size();iket++)
      // 	  {
      // 	    int iibra=pSystem_Tz->Groups[Orbitals[bra].GroupIndex][ibra];
      // 	    int iiket=pSystem_Tz->Groups[Orbitals[ket].GroupIndex][iket];
      // 	    val+=pSystem_Tz->get1B(A,iibra,iiket)*Orbitals[bra].Coeff[ibra]*Orbitals[ket].Coeff[iket];
      // 	  }
      val=(1.-1./A)*(1-beta)*get1B_p2(bra,ket)+beta*get1B_p2_plus_r2(bra,ket);
    }
  return val;
}

System_TzM_HF::DataType System_TzM_HF::get1B_p2(int bra,int ket) const
{
  DataType val=0;
  if(Orbitals[bra].isInAGroupWith(Orbitals[ket]) )
    {
      for(int ibra=0;ibra<Orbitals[bra].Coeff.size();ibra++)
	for(int iket=0;iket<Orbitals[ket].Coeff.size();iket++)
	  {
	    int iibra=pSystem_Tz->Groups[Orbitals[bra].GroupIndex][ibra];
	    int iiket=pSystem_Tz->Groups[Orbitals[ket].GroupIndex][iket];	    
	    val+=pSystem_Tz->get1B_p2(iibra,iiket)*Orbitals[bra].Coeff[ibra]*Orbitals[ket].Coeff[iket];
	  }
    }
  return val;
}
System_TzM_HF::DataType System_TzM_HF::get1B_p2_plus_r2(int bra,int ket) const
{
  DataType val=0;
  if(Orbitals[bra].isInAGroupWith(Orbitals[ket]) )
    {
      for(int ibra=0;ibra<Orbitals[bra].Coeff.size();ibra++)
	for(int iket=0;iket<Orbitals[ket].Coeff.size();iket++)
	  {
	    int iibra=pSystem_Tz->Groups[Orbitals[bra].GroupIndex][ibra];
	    int iiket=pSystem_Tz->Groups[Orbitals[ket].GroupIndex][iket];
	    if(iibra!=iiket) continue;
	    val+=pSystem_Tz->Orbitals[iibra].e*Orbitals[bra].Coeff[ibra]*Orbitals[ket].Coeff[iket];
	    //val+=pSystem_Tz->get1B_p2_plus_r2(iibra,iiket)*Orbitals[bra].Coeff[ibra]*Orbitals[ket].Coeff[iket];
	  }
    }
  return val;
}

System_TzM_HF::DataType System_TzM_HF::jplus(int bra,int ket) const
{
  DataType val=0;
  if(Orbitals[bra].par==Orbitals[ket].par && Orbitals[bra].tz==Orbitals[ket].tz &&
     (Orbitals[bra].mm-Orbitals[ket].mm)/2==1)
    {
      for(int ibra=0;ibra<Orbitals[bra].Coeff.size();ibra++)
	for(int iket=0;iket<Orbitals[ket].Coeff.size();iket++)
	  {
	    int iibra=pSystem_Tz->Groups[Orbitals[bra].GroupIndex][ibra];
	    int iiket=pSystem_Tz->Groups[Orbitals[ket].GroupIndex][iket];
	    val+=pSystem_Tz->jplus(iibra,iiket)*Orbitals[bra].Coeff[ibra]*Orbitals[ket].Coeff[iket];
	  }
    }
  return val;
} 

System_TzM_HF::DataType System_TzM_HF::jminus(int bra,int ket) const
{
  DataType val=0;
  if(Orbitals[bra].par==Orbitals[ket].par && Orbitals[bra].tz==Orbitals[ket].tz &&
     (Orbitals[ket].mm-Orbitals[bra].mm)/2==1)    
    {
      for(int ibra=0;ibra<Orbitals[bra].Coeff.size();ibra++)
	for(int iket=0;iket<Orbitals[ket].Coeff.size();iket++)
	  {
	    int iibra=pSystem_Tz->Groups[Orbitals[bra].GroupIndex][ibra];
	    int iiket=pSystem_Tz->Groups[Orbitals[ket].GroupIndex][iket];
	    val+=pSystem_Tz->jminus(iibra,iiket)*Orbitals[bra].Coeff[ibra]*Orbitals[ket].Coeff[iket];
	  }
    }
  return val;
}

System_TzM_HF::DataType System_TzM_HF::rL(int bra,int ket,int L) const
{
  DataType val=0;
  for(int ibra=0;ibra<Orbitals[bra].Coeff.size();ibra++)
    for(int iket=0;iket<Orbitals[ket].Coeff.size();iket++)
      {
	int iibra=pSystem_Tz->Groups[Orbitals[bra].GroupIndex][ibra];
	int iiket=pSystem_Tz->Groups[Orbitals[ket].GroupIndex][iket];
	val+=pSystem_Tz->rL(iibra,iiket,L)*Orbitals[bra].Coeff[ibra]*Orbitals[ket].Coeff[iket];
      }
  return val;
}
